# -*- coding: UTF-8 -*-
import os
import time
import rospy
import numpy as np
from sensor_msgs.msg import Imu
from sensor_msgs.msg import Image
import threading
from Queue import Queue
import matplotlib.pyplot as plt
import uuid

imuList=[]
imageList=Queue()
imuCheckingList=[]
imgProcessingList=Queue()

imuSendList=[]
imageSendList=Queue()

imu_buf=threading.Lock()
img_buf=threading.Lock()
data_train=None

def get_imu_record(path):
    list_files = []
    if os.path.exists(path):
        # 遍历文件夹，找到.npy文件
        for root, ds, fs in os.walk(path):
            for f in fs:
                if f.endswith('.npy'):
                    fullname = os.path.join(root, f)
                    list_files.append(fullname)
        for file in list_files:
            lab = os.path.split(os.path.split(file)[0])[-1]
    else:
        print('this path not exist')
    return list_files



def NORM(x):
    return np.math.sqrt(np.dot(x,x))

def dis_ACC(x, y):
    scale = 0.5
    diff = x - y
    dis = np.dot(x, y) / (NORM(x) * NORM(y) + 1e-8)
    dis = (1 - scale * dis) * NORM(diff)
    return dis

def estimate_dtw(slice, dis_func=dis_ACC):
    A=[]
    data=[]
    for item in slice:
        data.append(float(item.linear_acceleration.x))
        data.append(float(item.linear_acceleration.y))
        data.append(float(item.linear_acceleration.z))
        A.append(np.array(data))
        data=[]

    B=np.array(data_train)
    N_A = len(A)
    N_B = len(B)

    D = np.zeros([N_A, N_B])
    D[0, 0] = dis_func(A[0], B[0])

    #左边一列
    for i in range(1, N_A):
        D[i, 0] = D[i - 1, 0] + dis_func(A[i], B[0])
    #下边一行
    for j in range(1, N_B):
        D[0, j] = D[0, j - 1] + dis_func(A[0], B[j])
    #中间部分
    for i in range(1, N_A):
        for j in range(1, N_B):
            D[i, j] = dis_func(A[i], B[j]) + min(D[i - 1, j], D[i, j - 1], D[i - 1, j - 1])
    # 路径回溯
    i = N_A - 1
    j = N_B - 1
    count = 0
    d = np.zeros(max(N_A, N_B) * 3)
    path = []
    while True:
        if i > 0 and j > 0:
            path.append((i, j))
            m = min(D[i - 1, j], D[i, j - 1], D[i - 1, j - 1])
            if m == D[i - 1, j - 1]:
                d[count] = D[i, j] - D[i - 1, j - 1]
                i = i - 1
                j = j - 1
                count = count + 1

            elif m == D[i, j - 1]:
                d[count] = D[i, j] - D[i, j - 1]
                j = j - 1
                count = count + 1

            elif m == D[i - 1, j]:
                d[count] = D[i, j] - D[i - 1, j]
                i = i - 1
                count = count + 1

        elif i == 0 and j == 0:
            path.append((i, j))
            d[count] = D[i, j]
            count = count + 1
            break

        elif i == 0:
            path.append((i, j))
            d[count] = D[i, j] - D[i, j - 1]
            j = j - 1
            count = count + 1

        elif j == 0:
            path.append((i, j))
            d[count] = D[i, j] - D[i - 1, j]
            i = i - 1
            count = count + 1

    mean = np.sum(d) / count
    return mean, path[::-1], D

def handleIMU(data):
    imu_buf.acquire()
    imuList.append(data)
    imu_buf.release()

def handleImage(data):
    img_buf.acquire()
    imageList.put(data)
    img_buf.release()

def initialPublishImuByGroup():
    global imuList,imuSendList
    #print("sending imu ...")# -*- coding: UTF-8 -*-

    imu_buf.acquire()
    #print("imu size : "+str(len(imuList)))
    imuSendList=imuList
    imuList=[]
    imu_buf.release()
    imuSendList.reverse()
    while len(imuSendList)>0:
        imupub.publish(imuSendList.pop())
        imurate.sleep()
    #print("imu sended ! imu send list size : "+str(len(imuSendList)))


def initialPublishImgByGroup():
    global imageList,imageSendList
    #print("sending image ...")
    img_buf.acquire()
    #img_send_buf.acquire()
    #print("imgList size : "+str(imageList.qsize()))
    imageSendList=imageList
    imageList=Queue()
    img_buf.release()
    while imageSendList.qsize()>0:
        imgpub.publish(imageSendList.get())
        imgrate.sleep()
    #print("image sended ! image send list size : "+str(imageSendList.qsize()))
    #img_send_buf.release()

def imgProcess():
    global imageList,imgProcessingList
    img_buf.acquire()
    imgProcessingList=imageList
    imageList=Queue()
    img_buf.release()

    #print("image processed!")

processing=threading.Lock()

def imuProcessByDTW():
    global imuList,imuCheckingList
    imu_buf.acquire()
    #imu_check_buf.acquire()
    imuCheckingList=imuList
    imuList=[]
    imu_buf.release()
    processing.acquire()
    scores=estimate_dtw(imuCheckingList)
    if scores[0]>=0.7:
        imuCheckingList=data_replacing(imuCheckingList)
    processing.release()
    print("DTW done! score : "+str(scores[0]))
    #imu_check_buf.release()
    print("DTW done!")
    #return scores
sleep
def data_replacing(template):
    global i,data_train
    print("当前片段为："+str(i))
    print("检测到IMU数据异常，开始替换数据")
    for item,target in zip(template,data_train):
        item.linear_acceleration.x=target[0]
        item.linear_acceleration.y=target[1]
        item.linear_acceleration.z=target[2]
    print("替换完成")
    return template


def publishImuByGroup():
    global imuSendList, imuCheckingList
    imuSendList=imuCheckingList
    imuCheckingList=[]
    imuSendList.reverse()
    while len(imuSendList)>0:
        imupub.publish(imuSendList.pop())
        imurate.sleep()
    #print("imu sended just now!")

def publishImgByGroup():
    global imgProcessingList,imageSendList
    imageSendList=imgProcessingList
    imgProcessingList=Queue()
    while imageSendList.qsize()>0:
        imgpub.publish(imageSendList.get())
        imgrate.sleep()
    #print("image sended just now!")

i=0
def checkPublishStatus():
    global i
    FirstPublish=True
    while True:
        if FirstPublish:
            if len(imuList)>=600 and imageList.qsize()>=60:
                print("first public")
                imuPubTask=threading.Thread(target=initialPublishImuByGroup)
                imgPubTask=threading.Thread(target=initialPublishImgByGroup)

                imuPubTask.start()
                imgPubTask.start()

                imuPubTask.join()
                imgPubTask.join()
                FirstPublish=False
                i=i+1
                print("first public done")
        else:
            if  len(imuList)>=300  and imageList.qsize()>=30:

                imuCheckTask=threading.Thread(target=imuProcessByDTW)
                imgProcessTask=threading.Thread(target=imgProcess)

                imuCheckTask.start()
                imgProcessTask.start()

                imuCheckTask.join()
                imgProcessTask.join()
                print("imu & image processed")

                imuPubTask=threading.Thread(target=publishImuByGroup)
                imgPubTask=threading.Thread(target=publishImgByGroup)

                imuPubTask.start()
                imgPubTask.start()

                imuPubTask.join()
                imgPubTask.join()
                i=i+1
                print("imu & img publish done")

if __name__ == "__main__":

    rospy.init_node('FeatureProject',anonymous=True)
    imusub = rospy.Subscriber("/imu0", Imu,handleIMU, queue_size=20)
    imgsub=rospy.Subscriber("/cam0/image_raw",Image,handleImage,queue_size=20)

    imupub=rospy.Publisher("imuProcessed",Imu,queue_size=20)
    imgpub=rospy.Publisher("imageProcessed",Image,queue_size=20)
    imurate = rospy.Rate(600)
    imgrate = rospy.Rate(60)

    files = get_imu_record("../imu_template")
    for file in files:
        data_train = np.load(file)

    featureProject=threading.Thread(target=checkPublishStatus)
    featureProject.start()
    rospy.spin()

